# -*- coding: utf-8 -*-
import random
import sys
from PyQt4 import QtCore, QtGui
import os
import glob

try :
    from crowd_UI import crowd_ui
except :
    sys.path.append(r'C:\Users\onetera\Dropbox\R&D\crowd')
    from crowd_UI import crowd_ui

try:    
    import maya.OpenMaya as OM
    import maya.cmds as cmds
    import maya.OpenMayaUI as mui
    import maya.mel as mel
    import sip
    MAYA = 1
except :
    MAYA = 0

def getMayaWindow():
    if MAYA:
        ptr = mui.MQtUtil.mainWindow()
        return sip.wrapinstance(long(ptr), QtCore.QObject)
    else : return None

def ray( mesh , point ):
    sList = OM.MSelectionList()
    sList.add( mesh )
    item = OM.MDagPath()
    sList.getDagPath(0, item)
    item.extendToShape()
    fnMesh = OM.MFnMesh(item)    
    raySource = OM.MFloatPoint(-point[0], point[1], -point[2], -1.0)
    direction=(0.0, 1.0, 0.0)
    rayDir = OM.MFloatVector( direction[0], direction[1], direction[2] )
    faceIds = None
    triIds = None
    idsSorted = False
    testBothDirections = False
    worldSpace = OM.MSpace.kWorld
    maxParam = 999999
    accelParams = None
    sortHits = 0
    hitPoints = OM.MFloatPointArray()
    hitRayParams = OM.MFloatArray()
    hitFaces = OM.MIntArray()
    hitTris = None
    hitBarys1 = None
    hitBarys2 = None
    tolerance = 0.0001
    
    hit = fnMesh.allIntersections(raySource, rayDir, faceIds, triIds, idsSorted,
                worldSpace, maxParam, testBothDirections,
                accelParams, sortHits, hitPoints, hitRayParams,
                hitFaces, hitTris, hitBarys1, hitBarys2, tolerance)
    
    ff = hitPoints[0]
    return ff.x , ff.y , ff.z


def noise( obj , strength ):
    verts = cmds.polyEvaluate(obj , v=1 )   
    for n in range(verts) :
        x = random.uniform( strength*-1,strength )
        z = random.uniform( strength*-1,strength )
        cmds.move(x, 0, z, '%s.vtx[%d]' % ( obj , n), r=1)

def getCenter(pos1 , pos2):
    p1 = OM.MPoint( pos1[0] , pos1[1] , pos1[2] )
    p2 = OM.MPoint( pos2[0] , pos2[1] , pos2[2] )
    v1 = OM.MVector( p1 )
    v2 = OM.MVector( p2 )
    v3 = ( v1 + v2 ) / 2
    p3 = OM.MVector( v3 ) 
    return p3.x , p3.y , p3.z

def grdNoise(tg , pick = 7 , scale=13  , strength=5 ):
    verts = cmds.polyEvaluate( tg , v = 1 )
    for x in range( random.choice( range( 3,pick ) ) ):
        vert = random.choice( range(verts)  )
        cmds.select( tg + '.vtx[' + str(vert) + ']' )
        cmds.softSelect(tg, sse = 1 , ssd = scale , ssc = '0,1,2,1,0,2' ,
                        ssf = 2 )
        cmds.move( 0 , random.choice( range( strength) ) , 0 , r=1 )
    cmds.softSelect(sse = 0)    
    cmds.select( cl=1 )


class moneagi(crowd_ui):
    def __init__( self ):
        if getMayaWindow() != None:
            super( crowd_ui , self ).__init__(parent = getMayaWindow())
        else :
            super( crowd_ui , self ).__init__()        
        self.files = [ r"D:\WORK\Project\crowd\source\forward_walk01_pc.mb" ,
                       r"D:\WORK\Project\crowd\source\forward_walk02_pc.mb" ]
        self.grd = []
        self.scPlane = []
        self.modelList = []
        self.motionList = []
        self.scatterOBJ = []
        
    def setUI( self ):
        self.setupUi()
        self.tab1_crtGrd_btn.clicked.connect( self.createGround )
        self.tab1_noise_btn.clicked.connect( self.grdNoise )
        self.tab2_crtPlane_btn.clicked.connect( self.createScatterPlane )
        self.tab2_noise_btn.clicked.connect( self.scatterNoise )
        self.tab2_attach2GRD_btn.clicked.connect( self.attatchSPlane2grd )
        self.tab3_modelImport_btn.clicked.connect( self.modelImport )
        self.tab3_modelDelete_btn.clicked.connect( self.delMoldeList )
        self.tab3_motionImport_btn.clicked.connect( self.motionImport )
        self.tab3_motionDelete_btn.clicked.connect( self.delMotionList )
        self.tab4_scattering_btn.clicked.connect( self.scattering )
        self.tab4_delScat_btn.clicked.connect( self.deleteScatObj )
        self.tab4_animObj_btn.clicked.connect( self.animHeight )
        if not MAYA :
            self.tab4_scattering_btn.setEnabled( True )
        
    def createGround( self ):
        grdName = str( self.tab1_lineEdit.text() )
        grdWidth = self.tab1_width_spin.value()
        grdHeight = self.tab1_height_spin.value()
        grdsubWidth = self.tab1_subWidth_spin.value()
        grdsubHeight = self.tab1_subHeight_spin.value()
        if MAYA == 0 :            
            print 'Created Ground'
            return
        if grdName in cmds.ls( tr = 1 ):
            print "Ground is already existed."
            return
        cmds.undoInfo( openChunk = 1 )
        self.grd = cmds.polyPlane( n = grdName )
        cmds.setAttr( self.grd[1] + '.width' , grdWidth )
        cmds.setAttr( self.grd[1] + '.height' ,grdHeight )
        cmds.setAttr( self.grd[1] + '.subdivisionsHeight' , grdsubWidth )
        cmds.setAttr( self.grd[1] + '.subdivisionsWidth' , grdsubHeight )
        cmds.undoInfo( closeChunk = 1 )
        self.tab1_noise_btn.setEnabled( True )
        
    def grdNoise( self ):
        if self.grd == [] : return
        self.noisePick = self.tab1_noisePick_spin.value()
        self.noiseScale = self.tab1_noiseScale_spin.value()
        self.noiseStrength = self.tab1_noiseStrength_spin.value()
        cmds.undoInfo( openChunk=1 )
        grdNoise( self.grd[0] , pick = self.noisePick ,
                   scale = self.noiseScale  , strength = self.noiseStrength )
        cmds.undoInfo( closeChunk=1 )
        
    def createScatterPlane( self ):
        scPlaneName = str( self.tab1_lineEdit.text() )
        scatWidth = self.tab2_width_spin.value()
        scatHeight = self.tab2_height_spin.value()
        scatsubWidth = self.tab2_subWidth_spin.value()
        scatsubHeight = self.tab2_subHeight_spin.value()        

        for obj in self.scPlane:
            if cmds.objExists( obj ):
                print 'Scatter Plane is already existed.'                
                return
        if self.grd == '' :
            print 'Ground is not exist'
            return
        
        cmds.undoInfo( openChunk=1 )
        self.scPlane = cmds.polyPlane( n = scPlaneName )
        cmds.setAttr( self.scPlane[1] + '.width' ,scatWidth )
        cmds.setAttr( self.scPlane[1] + '.height' , scatHeight )
        cmds.setAttr( self.scPlane[1] + '.subdivisionsHeight' , scatsubWidth )
        cmds.setAttr( self.scPlane[1] + '.subdivisionsWidth' , scatsubHeight )        
        bbbox = cmds.xform( self.grd[0] , bb = 1 , q = 1 )
        cmds.setAttr( self.scPlane[0] +'.ty', bbbox[4] + bbbox[4]*0.1 )
        cmds.undoInfo( closeChunk=1 )
        self.tab2_attach2GRD_btn.setEnabled( True )
        self.tab4_scattering_btn.setEnabled( True )
        
        
    def scatterNoise( self ):
        scatNoise = self.tab2_noiseStrength_spin.value()
        if self.scPlane == [] :
            print 'Scatter Plane is not existed.'
            return
        cmds.undoInfo( openChunk=1 )
        noise( self.scPlane[0] , scatNoise )
        cmds.undoInfo( closeChunk=1 )
    
    def attatchSPlane2grd( self ):
        if self.scPlane == [] :
            print 'Scatter Plane is not existed.'
            return
        cmds.undoInfo( openChunk=1 )
        for avc in range( cmds.polyEvaluate( self.scPlane[0] , v=1 ) ):
            vert = self.scPlane[0] + '.vtx[' + str(avc) + ']'
            vert_pos = cmds.xform(vert , ws = 1, t=1 , q=1 )
            grd_pos = ray( self.grd[0] , vert_pos )
            cmds.xform( vert , ws = 1, t=grd_pos )
        cmds.undoInfo( closeChunk=1 )
    
    def modelImport( self ):
        modelList = QtGui.QFileDialog.getOpenFileNames(self, "Find Files",
                r'D:\WORK\Project\crowd\source' ,#QtCore.QDir.currentPath() ,
                'Maya (*.ma *.mb )' )
        self.modelList += [ str( t ) for t in list( modelList ) ]        
        if self.modelList :
            self.refreshModelList()
            
    def delMoldeList( self ):
        indexes = self.tab3_model_list.selectedIndexes()
        theIndexes = [ self.modelList[ ind.row() ]  for ind in indexes ]        
        for n in theIndexes:
            self.modelList.remove(n)
        self.refreshModelList()
        
    def refreshModelList( self ):        
        self.tab3_model_list.clear()
        for x in self.modelList:
            theItem = os.path.basename( x ).split('.')[0]
            print 'theItem ===> ',theItem
            self.tab3_model_list.addItem( theItem )
    
    def motionImport( self ):
        motionlList = QtGui.QFileDialog.getExistingDirectory(self, "Find Files",
                r'C:\Users\onetera\Documents\maya\projects\default\data' )
        self.motionList += [ str( motionlList ) ]        
        if self.motionList :
            self.refreshMotionList()
            
    def delMotionList( self ):
        indexes = self.tab3_motion_list.selectedIndexes()
        theIndexes = [ self.motionList[ ind.row() ]  for ind in indexes ]        
        for n in theIndexes:
            self.motionList.remove(n)
        self.refreshMotionList()
                    
    def refreshMotionList( self ):        
        self.tab3_motion_list.clear()
        for x in self.motionList:
            theItem = x.split('\\')[-1]
            self.tab3_motion_list.addItem( theItem )    

    def scattering( self ):
        if MAYA == 0 :            
            randStart = self.tab4_radStrt_spin.value()
            randEnd = self.tab4_radEnd_spin.value()
            randStartFrame = random.choice( range(randStart , randEnd ) )
            print self.tab3_model_list.items()
            return
        
        if self.grd == [] or self.scPlane == [] or self.modelList == [] or \
            self.motionList == [] or self.scatterOBJ != []:
            return
        self.scatterOBJ = []
        verts = cmds.polyEvaluate( self.scPlane[0] , v=1)        
        for vc in range( verts ):            
            modelFile = random.choice( self.modelList )           
            motionFiles = random.choice( self.motionList )
            scene_ls = [ x for x in cmds.ls(typ='transform') if 'geo' in x ]
            try : 
                cmds.file( modelFile , i=1  ,mnc = 0)
            except : 
                print vc
                continue
            new_ls = [ x for x in cmds.ls( typ='transform' ) if not x in scene_ls and 'geo' in x ]
            randStart = self.tab4_radStrt_spin.value()
            randEnd = self.tab4_radEnd_spin.value()
            randStartFrame = random.choice( range( randStart , randEnd ) )
            
            for i,n in enumerate( new_ls ):                
                theShape = cmds.listRelatives( n , c=1 )[0]
                switch = mel.eval( 'createHistorySwitch("%s",false)' % theShape )
                motionXML = [ x for x in glob.glob( motionFiles + '\\*.xml' ) if n in x ]
                print 'motionXML => ' , motionXML
                cacheNode = cmds.cacheFile(f=r"%s" % motionXML[0] ,
                         cnm='%sShape' % n, ia='%s.inp[0]' % switch ,attachFile=True)
                cmds.setAttr( '%s.playFromCache' % switch, 1 )                
                cmds.setAttr( cacheNode + '.startFrame' , randStartFrame )               

                new_item = cmds.rename( n , n+'_#' )        
                new_ls[i] = new_item    
            body_geo = [ x for x in new_ls if 'body' in x ]
            bb_pos = cmds.xform( body_geo[0] , bb=1 , q=1 )
            
            grp = cmds.createNode( 'transform' , n='CHR_grp_#'  )            
            bb_pos1 = [ bb_pos[0] , 0 , bb_pos[2] ]
            bb_pos2 = [ bb_pos[2] , 0 , bb_pos[5] ]
            center_pos = getCenter(bb_pos1 , bb_pos2)            
            cmds.setAttr(grp+'.t' , 0,0,center_pos[2] )
            cmds.parent( new_ls , grp )
            self.scatterOBJ.append( grp )
            vert_pos = cmds.xform( self.scPlane[0] + '.vtx['+str(vc)+']' , ws=1,t=1,q=1 )    
            cmds.xform( grp , ws=1,t=vert_pos )
        self.tab4_delScat_btn.setEnabled( True )
        self.tab4_animObj_btn.setEnabled( True )

    def deleteScatObj( self ):
        for x in self.scatterOBJ :            
            cmds.delete( x )
        self.tab4_animObj_btn.setEnabled( False )
        self.scatterOBJ = []
        
    def animHeight( self ):        
        for t in range(cmds.playbackOptions(min=1,q=1) , cmds.playbackOptions(max=1,q=1)) :
            cmds.currentTime(t)
            for body in self.scatterOBJ :
                bb_pos = cmds.xform( body , ws=1 , bb=1 , q=1 )
                bb_pos1 = [ bb_pos[0] , bb_pos[1] , bb_pos[2] ]
                bb_pos2 = [ bb_pos[3] , bb_pos[1] , bb_pos[5] ]
                center_pos = getCenter( bb_pos1 , bb_pos2 )
                #grp = cmds.listRelatives( body , p = 1) 
                ray_pos = ray( self.grd[0] , [ center_pos[0] , center_pos[1] + 100, center_pos[2] ] )
                cmds.setAttr(body + '.ty' ,ray_pos[1] )
                cmds.setKeyframe( body + '.ty' )

if __name__ == "__main__" :
    if MAYA :
        moneagi = moneagi()
        moneagi.setUI()
        moneagi.show()
    else :
        app = QtGui.QApplication(sys.argv)
        moneagi = moneagi()
        moneagi.setUI()
        moneagi.show()
        app.exec_()